from tools.common import run_async_script

import json
import os

from ioservice.common.mongo import insert_models
from ioservice.const import region_col
from models.administrative_division import RegionHierarchy, RegionType
from tools.region_data.alias_maping import *

current_dir = os.path.dirname(__file__)
data_dir = os.path.join(current_dir, "raw_data")

level2type = {
    1: RegionType.country,
    2: RegionType.province,
    3: RegionType.city,
    4: RegionType.district
}


def get_china_region_data():
    file_path = os.path.join(data_dir, "china_region.json")
    with open(file_path, mode='r', encoding='utf-8') as f:
        return json.load(f)


def get_special_region_data():
    file_path = os.path.join(data_dir, "special_regions.json")
    with open(file_path, mode='r', encoding='utf-8') as f:
        return json.load(f)


def get_foreign_regions_data():
    file_path = os.path.join(data_dir, "foreign_regions.json")
    with open(file_path, mode='r', encoding='utf-8') as f:
        return json.load(f)


def convert_china_region_to_model():
    def convert(path, level, data):
        r = RegionHierarchy(name=data['name'], path=path, type=level2type[level])
        if r.name == "市辖区":
            r.name = path.split(",")[-2]
        if 'code' in data:
            r.code = data['code']
        if r.name in country_alias:
            r.alias = [country_alias[r.name]]
        elif r.name in province_alias:
            r.alias = [province_alias[r.name]]
        results = [r]
        children = data.get('children', [])
        sub_path = "," if not path else path
        sub_path = sub_path + r.name + ","
        for child in children:
            results.extend(convert(sub_path, level+1, child))
        return results
    all_data = get_china_region_data()
    all_data.extend(get_special_region_data())
    china = {"name": "中国", "type": "country", "children": all_data}
    models = convert(None, 1, china)
    return models


def convert_foreign_regions_to_model():
    def convert(path, level, data):
        r = RegionHierarchy(name=data['n'], path=path, type=level2type[level])
        results = [r]
        children = data.get('s', [])
        sub_path = "," if not path else path
        sub_path = sub_path + r.name + ","
        for child in children:
            results.extend(convert(sub_path, level + 1, child))
        return results
    models = list()
    for item in get_foreign_regions_data():
        models.extend(convert(None, 1, item))
    return models


async def create_regions():
    models = convert_china_region_to_model()
    models.extend(convert_foreign_regions_to_model())
    result = await insert_models(region_col, models)
    print(result)


if __name__ == '__main__':
    run_async_script(create_regions())